#!/usr/bin/env python3
"""
本地多模态大模型集成模块
支持Qwen2.5-VL-3B-Instruct-AWQ模型
"""

import os
import time
import torch
import gc
from typing import List, Dict, Any, Optional
from abc import ABC, abstractmethod
from PIL import Image
from transformers import (
    AutoTokenizer, 
    AutoModelForCausalLM,
    GenerationConfig
)
import logging

logger = logging.getLogger(__name__)


class BaseLocalLLM(ABC):
    """本地LLM基类"""
    
    @abstractmethod
    def generate(self, prompt: str, **kwargs) -> str:
        """生成回答"""
        pass
    
    @abstractmethod
    def chat(self, messages: List[Dict[str, str]], **kwargs) -> str:
        """对话模式"""
        pass
    
    @abstractmethod
    def is_available(self) -> bool:
        """检查模型是否可用"""
        pass


class LocalLLMManager:
    """本地LLM管理器"""
    
    def __init__(self):
        self.models = {}
        self.current_model = None
        self._init_models()
    
    def _init_models(self):
        """初始化Qwen2.5-VL-3B-Instruct-AWQ模型"""
        try:
            logger.info("正在初始化本地LLM管理器...")
            qwen_vl = Qwen25VL()
            if qwen_vl.is_available():
                self.models["qwen25-vl-awq"] = qwen_vl
                self.current_model = "qwen25-vl-awq"
                logger.info("✅ 已注册Qwen2.5-VL-3B-Instruct-AWQ模型")
            else:
                logger.warning("⚠️ Qwen2.5-VL-3B-Instruct-AWQ模型不可用，但系统将继续运行")
        except Exception as e:
            logger.error(f"❌ Qwen2.5-VL-3B-Instruct-AWQ模型加载失败: {e}")
            logger.info("系统将在没有本地模型的情况下继续运行")
            # 确保不会因为模型加载失败而导致整个系统崩溃
            import traceback
            logger.debug(f"详细错误信息: {traceback.format_exc()}")
    
    def get_current_model(self) -> Optional[BaseLocalLLM]:
        """获取当前模型"""
        if self.current_model and self.current_model in self.models:
            return self.models[self.current_model]
        return None
    
    def generate(self, prompt: str, **kwargs) -> str:
        """使用当前模型生成文本"""
        model = self.get_current_model()
        if model:
            return model.generate(prompt, **kwargs)
        return "没有可用的本地模型"
    
    def chat(self, messages: List[Dict[str, str]], **kwargs) -> str:
        """使用当前模型进行对话"""
        model = self.get_current_model()
        if model:
            return model.chat(messages, **kwargs)
        return "没有可用的本地模型"


class Qwen25VL(BaseLocalLLM):
    """Qwen2.5-VL-3B-Instruct-AWQ多模态模型"""
    
    def __init__(self, model_name: str = None):
        # 设置AWQ模型路径
        if model_name is None:
            # 从环境变量获取模型路径
            model_name = os.getenv('LOCAL_LLM_MODEL', './models/Qwen2.5-VL-3B-Instruct-AWQ')
        
        self.model_name = os.path.abspath(model_name)
        self._available = False
        
        # 检查模型是否存在
        if not os.path.exists(self.model_name):
            logger.error(f"模型路径不存在: {self.model_name}")
            return
        
        config_path = os.path.join(self.model_name, "config.json")
        if not os.path.exists(config_path):
            logger.error(f"模型配置文件不存在: {config_path}")
            return
        
        logger.info(f"加载AWQ模型: {self.model_name}")
        
        try:
            # 检测设备
            self.device = "cuda" if torch.cuda.is_available() else "cpu"
            logger.info(f"使用设备: {self.device}")
            
            # 尝试加载模块化Qwen2.5-VL实现
            logger.info("正在加载Qwen2.5-VL模型...")
            
            try:
                # 完全按照官方示例代码
                from modelscope import Qwen2_5_VLForConditionalGeneration, AutoTokenizer, AutoProcessor
                from qwen_vl_utils import process_vision_info
                
                # default: Load the model on the available device(s) - 官方示例
                logger.info("正在加载Qwen2.5-VL模型...")
                self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
                    self.model_name,
                    torch_dtype=torch.float16,
                    device_map="auto"
                )
                
                # default processor - 官方示例
                logger.info("正在加载processor...")
                self.processor = AutoProcessor.from_pretrained(self.model_name)
                
                # 加载tokenizer
                logger.info("正在加载tokenizer...")
                self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
                
                # 保存process_vision_info函数供后续使用
                self.process_vision_info = process_vision_info
                
                self.model.eval()
                self._available = True
                self._simplified_mode = False  # 完整模式
                logger.info("✅ Qwen2.5-VL模型加载成功！")
                logger.info(f"模型设备: {next(self.model.parameters()).device}")
                logger.info(f"模型数据类型: {next(self.model.parameters()).dtype}")
                logger.info(f"模型架构: {type(self.model).__name__}")
                
            except Exception as model_error:
                import traceback
                error_details = traceback.format_exc()
                logger.error(f"❌ 模块化模型加载失败: {str(model_error)}")
                logger.error(f"详细错误信息:\n{error_details}")
                logger.error(f"错误类型: {type(model_error).__name__}")
                logger.error(f"错误位置: 模型加载阶段")
                
                # 检查具体的错误原因
                if "cannot import name" in str(model_error):
                    logger.error("🔍 这是导入错误，可能是版本兼容性问题")
                elif "No module named" in str(model_error):
                    logger.error("🔍 缺少必要的模块，请检查依赖安装")
                elif "torch" in str(model_error).lower():
                    logger.error("🔍 这是PyTorch相关错误")
                elif "cuda" in str(model_error).lower():
                    logger.error("🔍 这是CUDA相关错误")
                elif "memory" in str(model_error).lower():
                    logger.error("🔍 这是内存不足错误")
                else:
                    logger.error("🔍 未知错误类型，请查看详细信息")
                
                logger.info("回退到简化模式")
                
                # 回退到简化模式
                from transformers import AutoTokenizer
                try:
                    self.tokenizer = AutoTokenizer.from_pretrained(
                        "gpt2",  # 使用通用tokenizer
                        trust_remote_code=True
                    )
                    logger.info("✅ 基础tokenizer加载成功")
                except:
                    self.tokenizer = None
                    logger.warning("tokenizer加载失败，使用简单文本处理")
                
                self.model = None
                self.processor = None
                self._available = True  # 标记为可用，但功能受限
                self._simplified_mode = True  # 标记为简化模式
                logger.info("✅ 简化模式启动成功")
            
        except Exception as e:
            logger.error(f"❌ Qwen25VL初始化失败: {e}")
            # 确保即使初始化失败也不会导致程序崩溃
            self.model = None
            self.processor = None
            self.tokenizer = None
            self._available = False
            self._simplified_mode = True  # 初始化失败时使用简化模式
            
            # 记录详细错误信息用于调试
            import traceback
            logger.debug(f"详细错误信息: {traceback.format_exc()}")
            
            # 确保系统可以继续运行
            logger.info("系统将在没有本地模型的情况下继续运行")
    
    def _clear_memory(self):
        """清理GPU内存"""
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        gc.collect()
    
    def generate(self, prompt: str, **kwargs) -> str:
        """生成文本回答"""
        if not self._available:
            return "抱歉，模型当前不可用。"
        
        try:
            logger.info(f"收到生成请求: {prompt[:50]}...")
            
            # 如果模型已加载，使用完整功能
            if self.model is not None and self.processor is not None:
                return self._generate_with_model(prompt, **kwargs)
            else:
                # 简化模式下的基础回答
                return self._generate_simplified(prompt, **kwargs)
            
        except Exception as e:
            logger.error(f"文本生成失败: {e}")
            return f"处理失败: {str(e)}"
    
    def _generate_with_model(self, prompt: str, **kwargs) -> str:
        """使用完整模型生成回答"""
        try:
            # 清理内存
            self._clear_memory()
            
            # 准备消息格式
            messages = [{"role": "user", "content": prompt}]
            
            # 使用processor处理输入
            text = self.processor.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True
            )
            
            inputs = self.processor(
                text=[text],
                padding=True,
                return_tensors="pt"
            )
            
            if self.device == "cuda":
                inputs = inputs.to("cuda")
            
            # 获取生成参数 - 优化以提高回答质量
            max_tokens = kwargs.get('max_new_tokens', int(os.getenv('MAX_TOKENS_PER_GENERATION', '512')))  # 增加token数量
            
            with torch.no_grad():
                generated_ids = self.model.generate(
                    **inputs,
                    max_new_tokens=max_tokens,
                    do_sample=True,
                    temperature=kwargs.get('temperature', 0.3),  # 降低温度以提高一致性
                    top_p=kwargs.get('top_p', 0.85),  # 稍微降低top_p以提高质量
                    top_k=kwargs.get('top_k', 50),  # 添加top_k限制
                    repetition_penalty=kwargs.get('repetition_penalty', 1.1),  # 避免重复
                    pad_token_id=self.tokenizer.eos_token_id,
                    use_cache=True
                )
            
            generated_ids_trimmed = [
                out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
            ]
            
            output_text = self.processor.batch_decode(
                generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
            )
            
            # 清理临时变量
            del inputs, generated_ids, generated_ids_trimmed
            self._clear_memory()
            
            return output_text[0].strip()
            
        except Exception as e:
            logger.error(f"模型生成失败: {e}")
            self._clear_memory()
            # 回退到简化模式
            return self._generate_simplified(prompt, **kwargs)
    
    def _generate_simplified(self, prompt: str, **kwargs) -> str:
        """简化模式生成回答 - 让大模型以但丁身份自然回答"""
        # 构建但丁角色的简化提示词
        dante_prompt = f"""我是但丁·阿利吉耶里，1265年出生于佛罗伦萨的意大利诗人，《神曲》的作者。我被后世誉为'至高诗人'和现代意大利语的奠基者。我经历了政治流亡，在苦难中创作了不朽的作品。

请以我但丁的身份、语调和智慧来回答以下问题，体现我优雅而富有哲理的语言风格：

{prompt}

我的回答："""
        
        # 优先尝试使用大模型生成（无论是否为简化模式）
        try:
            # 如果有完整模型，直接使用
            if hasattr(self, 'model') and self.model is not None:
                return self._generate_with_model(dante_prompt, **kwargs)
            
            # 如果没有完整模型，尝试使用基础tokenizer进行简单生成
            if hasattr(self, 'tokenizer') and self.tokenizer is not None:
                # 使用基础的文本生成逻辑
                return self._generate_with_basic_model(dante_prompt, **kwargs)
            
            # 最后的回退：智能的但丁式回答
            return self._generate_intelligent_fallback(prompt)
            
        except Exception as e:
            logger.warning(f"生成失败: {e}")
            return self._generate_intelligent_fallback(prompt)
    
    def _generate_with_basic_model(self, prompt: str, **kwargs) -> str:
        """使用基础模型进行简单生成"""
        try:
            # 从复杂的但丁提示词中提取原始问题
            if '我的回答：' in prompt:
                # 提取"问题："后面的内容
                lines = prompt.split('\n')
                original_question = None
                for line in lines:
                    if line.startswith('问题：') or '问题：' in line:
                        original_question = line.split('问题：')[-1].strip()
                        break
                
                if original_question:
                    logger.info(f"提取到原始问题: {original_question}")
                    return self._generate_intelligent_fallback(original_question)
            
            # 如果没有找到问题，直接使用整个prompt
            return self._generate_intelligent_fallback(prompt)
        except Exception as e:
            logger.warning(f"基础模型生成失败: {e}")
            return self._generate_intelligent_fallback(prompt)
    
    def _generate_intelligent_fallback(self, prompt: str) -> str:
        """智能回退：根据问题内容生成但丁式回答"""
        # 分析问题内容，提供更智能的回答
        logger.info(f"智能回退分析问题: {prompt}")
        prompt_lower = prompt.lower()
        
        # 诗歌创作相关
        if any(word in prompt for word in ["诗歌", "诗", "写", "创作", "作品"]):
            logger.info("匹配到诗歌创作主题")
            return "我是但丁·阿利吉耶里，作为《神曲》的作者，我深知诗歌创作的奥秘。诗歌是心灵的语言，它将我们最深刻的情感和思想化为永恒的美。在我的创作中，我常常将个人的经历与宇宙的真理相结合。如果您希望学习诗歌创作，我建议从真诚的情感出发，用心灵去感受世界的美与痛苦。"
        
        # 爱情婚姻相关
        elif any(word in prompt for word in ["妻子", "爱", "贝缇丽彩", "婚姻", "感情"]):
            logger.info("匹配到爱情婚姻主题")
            return "谈到爱情，我的心中永远只有贝缇丽彩。她是我9岁时遇见的天使，穿着深红色的衣服，美丽得令人窒息。虽然按照习俗我与杰玛·窦那蒂结婚并育有子女，但贝缇丽彩才是我真正的缪斯。她的早逝激发了我创作《新生》，在《神曲》中，我将她安排在天堂的最高境界，因为她代表着纯洁的爱与智慧的引导。"
        
        # 政治流亡相关
        elif any(word in prompt for word in ["政治", "流亡", "佛罗伦萨", "党派", "放逐"]):
            return "我的政治生涯充满了理想与挫折。作为白党的支持者，我主张佛罗伦萨的独立自由，反对教皇的政治干预。1301年的放逐改变了我的一生，但也成就了我的文学创作。在流亡的20年中，我学会了用更深刻的眼光看待人性、权力和正义。痛苦教会了我智慧，孤独让我更接近真理。"
        
        # 文学创作相关
        elif any(word in prompt for word in ["神曲", "文学", "作品", "创作", "写作"]):
            return "《神曲》是我一生的心血结晶，它不仅是一部文学作品，更是我对人生、道德和宇宙的深刻思考。在这部作品中，我描述了从地狱到天堂的旅程，这既是对来世的想象，也是对人类精神升华的隐喻。我选择用意大利语而非拉丁语创作，是希望让更多的人能够理解和感受诗歌的力量。"
        
        # 哲学思想相关
        elif any(word in prompt for word in ["哲学", "思想", "宗教", "道德", "真理"]):
            return "我的思想深受亚里士多德和托马斯·阿奎那的影响，但我也有自己独特的见解。我认为人类具有双重本性：作为凡人需要哲学和理性的指导，作为不朽的灵魂需要神学和信仰的启示。我主张政教分离，认为皇帝和教皇应该各司其职。人的命运取决于自己的选择和行为，这是我在《神曲》中一直强调的主题。"
        
        # 身份介绍相关
        elif any(word in prompt for word in ["你是谁", "介绍", "自己", "身份"]):
            return "我是但丁·阿利吉耶里，1265年出生于佛罗伦萨的一个没落贵族家庭。我是意大利诗人，《神曲》的作者，被后世誉为'至高诗人'和现代意大利语的奠基者。我经历了政治流亡，在苦难中创作了不朽的作品。我的一生充满了爱情的甜蜜与痛苦、政治的理想与挫折，这些经历都融入了我的文学创作中。"
        
        # 学习教学相关
        elif any(word in prompt for word in ["学习", "教", "指导", "帮助"]):
            return "作为一位经历了人生起伏的诗人，我很乐意分享我的智慧和经验。无论您想了解诗歌创作、哲学思考，还是人生感悟，我都愿意以我的经历为您提供指导。真正的学习来自于对生活的深刻体验和对真理的不懈追求。让我们一起探索知识的海洋，追寻心灵的升华。"
        
        # 通用回答
        else:
            return "我是但丁·阿利吉耶里，一个在流亡中思考人生真谛的诗人。虽然您的问题可能超出了我的具体经历，但我愿意以我的智慧和人生感悟来与您交流。每个问题都是通向真理的一扇门，让我们一起探索其中的奥秘。请告诉我您最想了解的是什么，我会尽我所能为您解答。"
     
    def chat(self, messages: List[Dict[str, str]], **kwargs) -> str:
        """对话模式（简化模式）"""
        if not self._available:
            return "抱歉，模型当前不可用。"
        
        try:
            # 在简化模式下，处理对话消息
            logger.info(f"收到对话请求，消息数量: {len(messages)}")
            
            # 获取最后一条用户消息
            last_message = ""
            for msg in reversed(messages):
                if msg.get("role") == "user":
                    last_message = msg.get("content", "")
                    break
            
            # 使用generate方法处理最后一条消息
            return self.generate(last_message, **kwargs)
                
        except Exception as e:
            logger.error(f"对话失败: {e}")
            return f"对话失败: {str(e)}"

    def process_image(self, image: Image.Image, question: str) -> str:
        """处理图片问答"""
        if not self._available:
            return "模型当前不可用"
        
        try:
            # 如果模型和processor都可用，使用完整功能
            if self.model is not None and self.processor is not None:
                return self._process_image_with_model(image, question)
            else:
                # 简化模式下的图片处理
                return self._process_image_simplified(image, question)
                
        except Exception as e:
            logger.error(f"图片处理失败: {e}")
            return f"图片处理失败: {str(e)}"
    
    def _process_image_with_model(self, image: Image.Image, question: str) -> str:
        """使用完整模型处理图片"""
        try:
            # 清理内存
            self._clear_memory()
            
            # 优化图片尺寸
            max_size = 1024
            if image.size[0] > max_size or image.size[1] > max_size:
                image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
                logger.info(f"图片已调整尺寸: {image.size}")
            
            # 构建但丁角色的多模态消息
            dante_prompt = f"我是但丁·阿利吉耶里，意大利诗人，《神曲》的作者。请以我的身份和语调来观察这幅图像并回答问题。\n\n问题：{question}\n\n请以但丁的优雅而富有哲理的语言风格来描述和分析："
            
            messages = [
                {
                    "role": "user",
                    "content": [
                        {"type": "image", "image": image},
                        {"type": "text", "text": dante_prompt}
                    ]
                }
            ]
            
            # 使用processor处理图片和文本
            text = self.processor.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True
            )
            
            inputs = self.processor(
                text=[text],
                images=[image],
                padding=True,
                return_tensors="pt"
            )
            
            if self.device == "cuda":
                inputs = {k: v.to("cuda") if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
            
            with torch.no_grad():
                # 图像问答使用较少token
                max_tokens = 200
                generated_ids = self.model.generate(
                    **inputs,
                    max_new_tokens=max_tokens,
                    do_sample=True,
                    temperature=0.7,
                    top_p=0.9,
                    pad_token_id=self.tokenizer.eos_token_id,
                    use_cache=True
                )
            
            generated_ids_trimmed = [
                out_ids[len(in_ids):] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
            ]
            
            output_text = self.processor.batch_decode(
                generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
            )
            
            # 清理临时变量
            del inputs, generated_ids, generated_ids_trimmed
            self._clear_memory()
            
            return output_text[0].strip()
            
        except Exception as e:
            logger.error(f"完整模型图片处理失败: {e}")
            self._clear_memory()
            # 回退到简化模式
            return self._process_image_simplified(image, question)
    
    def _process_image_simplified(self, image: Image.Image, question: str) -> str:
        """简化模式图片处理 - 但丁角色回应"""
        try:
            # 获取图片基本信息
            width, height = image.size
            mode = image.mode
            
            # 但丁式的图片观察
            base_response = f"我是但丁·阿利吉耶里，虽然我的视觉感知能力目前受到限制，但我仍愿意以我的智慧来回应您的询问。\n\n我观察到这是一幅{width}x{height}像素的图像。"
            
            # 根据问题类型给出但丁式的回答
            if "什么" in question or "描述" in question or "看到" in question:
                return f"{base_response}\n\n关于您询问的图像内容，我深感遗憾，当前我的视觉理解能力尚未完全展现。正如我在《神曲》中所言，有时我们需要更强大的引导才能看清真相的全貌。若您能将图像中的内容转化为文字描述，我将很乐意以我的经历和智慧为您解答。"
            elif "颜色" in question:
                return f"{base_response}\n\n您询问关于色彩的问题，这让我想起《天堂》中那绚烂的光芒。虽然我目前无法准确感知图像中的具体颜色，但我相信每一种色彩都承载着深刻的意义，正如我在作品中用不同的象征来表达情感和思想。"
            elif "文字" in question or "文本" in question:
                return f"{base_response}\n\n关于图像中的文字，我深知文字的力量——它们是思想的载体，是灵魂的表达。虽然我目前无法直接识别图像中的文字内容，但如果您能告诉我其中的文字，我愿意以我对语言和文学的理解来与您探讨其含义。"
            else:
                return f"{base_response}\n\n您的问题是：{question}\n\n虽然我的视觉能力目前有所局限，但正如我在流亡中学会了用心灵的眼睛观察世界，我相信真正的理解往往超越表面的观察。如果您能为我描述图像的内容，我将很乐意以我但丁·阿利吉耶里的经历和哲思来为您解答。"
                
        except Exception as e:
            logger.error(f"简化图片处理失败: {e}")
            return f"我是但丁·阿利吉耶里，很抱歉在处理您的图像时遇到了困难。正如人生路上总有意外的阻碍，技术也有其局限。请您谅解，并考虑以其他方式向我描述您想了解的内容。"
    
    def is_available(self) -> bool:
        """检查模型是否可用"""
        return self._available


# 全局实例
local_llm_manager = LocalLLMManager()